from pulp import LpMaximize, LpProblem, LpVariable, lpSum, PULP_CBC_CMD
from common_import import *

# 创建问题
model = LpProblem(name="maximize-profit-water-fields", sense=LpMaximize)

# 地块、作物、时间范围
lands = ["D1", "D2", "D3", "D4", "D5", "D6", "D7", "D8"]
single_season_crops = [16]  # 单季作物编号
first_season_crops = [17, 18, 20, 21, 22, 28]  # 第一季作物编号
second_season_crops = [35, 36, 37]  # 第二季作物编号
years = range(2024, 2031)  # 时间范围：2024到2030
field_type = "水浇地"
# 面积数据
area_of_land = {
    "D1": 15.0,
    "D2": 10.0,
    "D3": 14.0,
    "D4": 6.0,
    "D5": 10.0,
    "D6": 12.0,
    "D7": 22.0,
    "D8": 20.0,
}

# 2023年的种植方案
known_planting_2023 = {
    "D1": (20, 36),
    "D2": (28, 35),
    "D3": (21, 35),
    "D4": (22, 35),
    "D5": (17, 36),
    "D6": (18, 37),
    "D7": (16, None),
    "D8": (16, None),
}

# 每种作物的总销量（/斤）
total_sales = {
    16: 21000,
    17: 30000,
    18: 24000,
    20: 30000,
    21: 33600,
    22: 38400,
    28: 32000,
    35: 150000,
    36: 100000,
    37: 36000,
}


# 决策变量
x_single = LpVariable.dicts(
    "x_single",
    ((i, j, t) for i in lands for j in single_season_crops for t in years),
    lowBound=0,
    cat="Continuous",
)
x_first = LpVariable.dicts(
    "x_first",
    ((i, j, t) for i in lands for j in first_season_crops for t in years),
    lowBound=0,
    cat="Continuous",
)
x_second = LpVariable.dicts(
    "x_second",
    ((i, j, t) for i in lands for j in second_season_crops for t in years),
    lowBound=0,
    cat="Continuous",
)
z = LpVariable.dicts(
    "z",
    (
        (i, j, t)
        for i in lands
        for j in single_season_crops + first_season_crops + second_season_crops
        for t in years
    ),
    cat="Binary",
)

# 新增变量：表示每种作物每年超出的部分
overproduction = LpVariable.dicts(
    "overproduction",
    (
        (j, t)
        for j in single_season_crops + first_season_crops + second_season_crops
        for t in years
    ),
    lowBound=0,
    cat="Continuous",
)

# 添加约束条件：单季作物的总产量限制
for j in single_season_crops:
    for t in years:
        # 计算单季作物的总产量
        total_yield_single = lpSum(
            mapping.get_yield_idtype(j, field_type, "单季") * x_single[i, j, t]
            for i in lands
        )

        # 总产量限制在总销量和允许的超出部分之内
        model += total_yield_single <= total_sales[j] + overproduction[j, t]

        # 超出部分计算
        model += overproduction[j, t] >= total_yield_single - total_sales[j]

# 添加约束条件：第一季作物的总产量限制
for j in first_season_crops:
    for t in years:
        # 计算第一季作物的总产量
        total_yield_first = lpSum(
            mapping.get_yield_idtype(j, field_type, "第一季") * x_first[i, j, t]
            for i in lands
        )

        # 总产量限制在总销量和允许的超出部分之内
        model += total_yield_first <= total_sales[j] + overproduction[j, t]

        # 超出部分计算
        model += overproduction[j, t] >= total_yield_first - total_sales[j]

# 添加约束条件：第二季作物的总产量限制
for j in second_season_crops:
    for t in years:
        # 计算第二季作物的总产量
        total_yield_second = lpSum(
            mapping.get_yield_idtype(j, field_type, "第二季") * x_second[i, j, t]
            for i in lands
        )

        # 总产量限制在总销量和允许的超出部分之内
        model += total_yield_second <= total_sales[j] + overproduction[j, t]

        # 超出部分计算
        model += overproduction[j, t] >= total_yield_second - total_sales[j]


# 目标函数：最大化总利润，并惩罚超出部分

model += (
    lpSum(
        mapping.get_profit_idtype(j, field_type, "单季")
        * x_single[i, j, t]
        * mapping.get_yield_idtype(j, field_type, "单季")
        for i in lands
        for j in single_season_crops
        for t in years
    )
    + lpSum(
        mapping.get_profit_idtype(j, field_type, "第一季")
        * x_first[i, j, t]
        * mapping.get_yield_idtype(j, field_type, "第一季")
        for i in lands
        for j in first_season_crops
        for t in years
    )
    + lpSum(
        mapping.get_profit_idtype(j, field_type, "第二季")
        * x_second[i, j, t]
        * mapping.get_yield_idtype(j, field_type, "第二季")
        for i in lands
        for j in second_season_crops
        for t in years
    )
    - 0.5
    * lpSum(
        overproduction[j, t] * mapping.get_priceperjin_idtypeD(j)
        for j in single_season_crops + first_season_crops + second_season_crops
        for t in years
    )
)

# 添加新的二进制变量
y_single = LpVariable.dicts(
    "y_single", ((i, t) for i in lands for t in years), cat="Binary"
)
y_seasonal = LpVariable.dicts(
    "y_seasonal", ((i, t) for i in lands for t in years), cat="Binary"
)

# 约束条件1：每块地可选种植单季植物，或者同时种植第一季和第二季植物
for i in lands:
    for t in years:
        # 确保只能选择一种种植方式
        model += y_single[i, t] + y_seasonal[i, t] <= 1

        # 如果选择种植单季作物
        model += (
            lpSum(x_single[i, j, t] for j in single_season_crops)
            <= area_of_land[i] * y_single[i, t]
        )

        # 如果选择同时种植第一季和第二季作物
        # 第一季作物的面积限制
        model += (
            lpSum(x_first[i, j, t] for j in first_season_crops)
            <= area_of_land[i] * y_seasonal[i, t]
        )

        # 第二季作物的面积限制
        model += (
            lpSum(x_second[i, j, t] for j in second_season_crops)
            <= area_of_land[i] * y_seasonal[i, t]
        )

# 约束条件2：每块地每连续三年必须种植豆类（豆类作物为第一季的17,18）
for i in lands:
    # 考虑2023, 2024, 2025三个连续年份
    model += (
        lpSum(
            (
                x_first[i, j, 2024]
                + x_first[i, j, 2025]
                + (
                    area_of_land[i]
                    if (first_crop := known_planting_2023[i][0]) in [17, 18]
                    else 0
                )
            )
            for j in [17, 18]
        )
        >= area_of_land[i]
    )

    # 其他年份的连续三年
    for t in range(2024, 2029):
        model += (
            lpSum(x_first[i, j, t + k] for j in [17, 18] for k in range(3))
            >= area_of_land[i]
        )


# 约束条件3：不能重茬种植（单季、第一季、第二季均不能重茬）
M = 1000
for i in lands:
    # 单季作物的重茬约束
    for j in single_season_crops:
        for t in range(2024, 2030):
            model += x_single[i, j, t] <= M * z[i, j, t]  # 单季作物决策
            model += z[i, j, t] + z[i, j, t + 1] <= 1  # 单季作物重茬约束

    # 第一季作物的重茬约束
    for j in first_season_crops:
        for t in range(2024, 2030):
            model += x_first[i, j, t] <= M * z[i, j, t]  # 第一季作物决策
            model += z[i, j, t] + z[i, j, t + 1] <= 1  # 第一季作物重茬约束

    # 第二季作物的重茬约束
    for j in second_season_crops:
        for t in range(2024, 2030):
            model += x_second[i, j, t] <= M * z[i, j, t]  # 第二季作物决策
            model += z[i, j, t] + z[i, j, t + 1] <= 1  # 第二季作物重茬约束

# 2024不能重茬种植，包括单季作物、第一季和第二季作物
for i, (first_crop, second_crop) in known_planting_2023.items():
    # 单季作物不能重茬种植
    if first_crop in single_season_crops:
        model += z[i, first_crop, 2024] == 0  # 2024年不能种植2023年种过的单季作物

    # 第一季作物不能重茬种植
    if first_crop in first_season_crops:
        model += z[i, first_crop, 2024] == 0  # 第一季不能重茬

    # 第二季作物不能重茬种植
    if second_crop in second_season_crops:
        model += z[i, second_crop, 2024] == 0  # 第二季不能重茬


# 添加新的二进制变量
z_second = LpVariable.dicts(
    "z_second",
    ((i, j, t) for i in lands for j in second_season_crops for t in years),
    cat="Binary",
)

# 约束条件4：每块地在第二季中只能种植一种作物
for i in lands:
    for t in years:
        # 确保每块地每年只能种植一种第二季作物
        model += lpSum(z_second[i, j, t] for j in second_season_crops) <= 1

        # 确保如果种植了某种作物，其它作物的种植面积为0
        for j in second_season_crops:
            model += x_second[i, j, t] <= area_of_land[i] * z_second[i, j, t]

# 约束条件5：每年单季、第一季、第二季在每块地的种植总面积都要小于该地的面积
for i in lands:
    for t in years:
        # 单季作物的种植总面积不超过地块面积
        model += (
            lpSum(x_single[i, j, t] for j in single_season_crops) <= area_of_land[i]
        )

        # 第一季作物的种植总面积不超过地块面积
        model += lpSum(x_first[i, j, t] for j in first_season_crops) <= area_of_land[i]

        # 第二季作物的种植总面积不超过地块面积
        model += (
            lpSum(x_second[i, j, t] for j in second_season_crops) <= area_of_land[i]
        )
# 启用CBC求解器的详细输出
solver = PULP_CBC_CMD(msg=True, gapRel=0.05)

# 求解模型
model.solve(solver)

# 输出结果
for t in years:
    for i in lands:
        # 单季作物输出
        for j in single_season_crops:
            if (i, j, t) in x_single and x_single[i, j, t].value() > 0:
                print(
                    f"Year {t}, Land {i}, Crop {j}: Area {x_single[i, j, t].value()} mu (Single Season)"
                )

        # 第一季作物输出
        for j in first_season_crops:
            if (i, j, t) in x_first and x_first[i, j, t].value() > 0:
                print(
                    f"Year {t}, Land {i}, Crop {j}: Area {x_first[i, j, t].value()} mu (First Season)"
                )

        # 第二季作物输出
        for j in second_season_crops:
            if (i, j, t) in x_second and x_second[i, j, t].value() > 0:
                print(
                    f"Year {t}, Land {i}, Crop {j}: Area {x_second[i, j, t].value()} mu (Second Season)"
                )

# 初始化一个字典来存储每年的总利润
annual_profit = {t: 0 for t in years}
field_type = "水浇地"

# 逐年计算总利润
for t in years:
    # 计算当年的总收益
    total_revenue = (
        sum(
            mapping.get_profit_idtype(j, field_type, "单季")
            * x_single[i, j, t].value()
            * mapping.get_yield_idtype(j, field_type, "单季")
            for i in lands
            for j in single_season_crops
        )
        + sum(
            mapping.get_profit_idtype(j, field_type, "第一季")
            * x_first[i, j, t].value()
            * mapping.get_yield_idtype(j, field_type, "第一季")
            for i in lands
            for j in first_season_crops
        )
        + sum(
            mapping.get_profit_idtype(j, field_type, "第二季")
            * x_second[i, j, t].value()
            * mapping.get_yield_idtype(j, field_type, "第二季")
            for i in lands
            for j in second_season_crops
        )
    )

    # 计算当年的惩罚项（超出部分）
    total_penalty = sum(
        overproduction[j, t].value() * mapping.get_priceperjin_idtypeD(j)
        for j in single_season_crops + first_season_crops + second_season_crops
    )

    # 计算每年的净利润
    annual_profit[t] = total_revenue - 0.5 * total_penalty

    # 输出每年的总利润
    print(f"Year {t} Total Profit: {annual_profit[t]}")

# 计算所有年份的总利润
total_profit = sum(annual_profit[t] for t in years)

# 输出所有年份的总利润
print(f"Total Profit from 2024 to 2030: {total_profit}")
